package com.chatplus.application.processor.listener;

import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.processor.constant.AIConstants;
import com.chatplus.application.processor.handler.ChatAiUtilHandler;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.service.v3.ModelEventSourceListener;
import com.zhipu.oapi.service.v3.SseMeta;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;

/**
 * ChatGLM监听器
 *
 * @author chj
 * @date 2024/1/4
 **/
public class ChatGLMEventSourceListener extends ModelEventSourceListener {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatGLMEventSourceListener.class);
    String reply = "";

    private SseMeta sseMeta;
    protected final ChatAiUtilHandler chatAiUtilHandler;

    public ChatGLMEventSourceListener(WebSocketSession session, String platform, String prompt) {
        chatAiUtilHandler = new ChatAiUtilHandler(session, platform, prompt);
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        LOGGER.message("连接成功").context("eventSource", eventSource.request()).context("response", response).info();
        chatAiUtilHandler.sendMessage(AIConstants.START_LABEL);
    }

    @Override
    public void onClosed(EventSource eventSource) {
        LOGGER.message("关闭连接").warn();
        eventSource.cancel();
    }

    @Override
    public void onFailure(@NotNull EventSource eventSource, Throwable t, Response response) {
        if (Objects.isNull(response)) {
            return;
        }
        try {
            ResponseBody body = response.body();
            if (Objects.nonNull(body)) {
                String bodyStr = body.string();
                LOGGER.message("连接异常").context("data", bodyStr).exception(t).error();
                chatAiUtilHandler.replyErrorMessage("服务端发生异常，等待管理员修复,第三方返回: " + bodyStr);

            } else {
                LOGGER.message("连接异常").context("data", response).exception(t).error();
                chatAiUtilHandler.replyErrorMessage("服务端发生异常，等待管理员修复!");
            }
        } catch (IOException e) {
            LOGGER.message("连接异常").exception(e).error();
        }
        eventSource.cancel();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {
        LOGGER.message("ChatGLM返回数据").context("data", data).info();
        ObjectMapper objectMapper = SpringUtil.getBean(ObjectMapper.class);
        JSONObject jsonObject = JSONUtil.parseObj(data);
        String content = jsonObject.getStr("data");
        try {
            if (StringUtils.isNotEmpty(content)) {
                reply += content;
                chatAiUtilHandler.sendMessage(content);
            }
            if ("finish".equals(type)) {
                String meta = jsonObject.getStr("meta");
                this.sseMeta = objectMapper.readValue(meta, SseMeta.class);
                LOGGER.message("回答完成").info();
                chatAiUtilHandler.saveChatHistory(reply);
                chatAiUtilHandler.sendMessage(AIConstants.END_LABEL);
            }
        } catch (JsonProcessingException e) {
            chatAiUtilHandler.replyErrorMessage("服务端发生异常，等待管理员修复!");
            LOGGER.message("JSON 解析错误").exception(e).error();
        }
    }

    @Override
    public String getOutputText() {
        return reply;
    }

    @Override
    public CountDownLatch getCountDownLatch() {
        return null;
    }

    @Override
    public SseMeta getMeta() {
        return sseMeta;
    }
}